import carla
import random
import os
from queue import Queue
from queue import Empty

import pygame

from manual_control import HUD, KeyboardControl
from manual_control import CameraManager, get_actor_display_name, find_weather_presets
from manual_control import CollisionSensor, LaneInvasionSensor, GnssSensor, IMUSensor
import sys


class World(object):
    def __init__(self, carla_world, hud, player, args=None):
        self.world = carla_world
        self.sync = True
        self.actor_role_name = 'no'
        try:
            self.map = self.world.get_map()
        except RuntimeError as error:
            print('RuntimeError: {}'.format(error))
            print('  The server could not send the OpenDRIVE (.xodr) file:')
            print(
                '  Make sure it exists, has the same name of your town, and is correct.'
            )
            sys.exit(1)
        self.hud = hud
        self.player = player
        self.collision_sensor = None
        self.lane_invasion_sensor = None
        self.gnss_sensor = None
        self.imu_sensor = None
        self.radar_sensor = None
        self.camera_manager = None
        self._weather_presets = find_weather_presets()
        self._weather_index = 0
        self._actor_filter = None
        self._actor_generation = None
        self._gamma = 2.2
        self.restart()
        self.world.on_tick(hud.on_world_tick)
        self.recording_enabled = False
        self.recording_start = 0
        self.constant_velocity_enabled = False
        self.show_vehicle_telemetry = False
        self.current_map_layer = 0
        self.map_layer_names = [
            carla.MapLayer.NONE,
            carla.MapLayer.Buildings,
            carla.MapLayer.Decals,
            carla.MapLayer.Foliage,
            carla.MapLayer.Ground,
            carla.MapLayer.ParkedVehicles,
            carla.MapLayer.Particles,
            carla.MapLayer.Props,
            carla.MapLayer.StreetLights,
            carla.MapLayer.Walls,
            carla.MapLayer.All,
        ]

    def restart(self):
        self.player_max_speed = 1.589
        self.player_max_speed_fast = 3.713
        # Keep same camera config if the camera manager exists.
        cam_index = self.camera_manager.index if self.camera_manager is not None else 0
        cam_pos_index = (
            self.camera_manager.transform_index
            if self.camera_manager is not None
            else 0
        )

        self.collision_sensor = CollisionSensor(self.player, self.hud)
        self.lane_invasion_sensor = LaneInvasionSensor(self.player, self.hud)
        self.gnss_sensor = GnssSensor(self.player)
        self.imu_sensor = IMUSensor(self.player)
        self.camera_manager = CameraManager(self.player, self.hud, self._gamma)
        self.camera_manager.transform_index = cam_pos_index
        self.camera_manager.set_sensor(cam_index, notify=False)
        actor_type = get_actor_display_name(self.player)
        self.hud.notification(actor_type)

        if self.sync:
            self.world.tick()
        else:
            self.world.wait_for_tick()

    def next_weather(self, reverse=False):
        self._weather_index += -1 if reverse else 1
        self._weather_index %= len(self._weather_presets)
        preset = self._weather_presets[self._weather_index]
        self.hud.notification('Weather: %s' % preset[1])
        self.player.get_world().set_weather(preset[0])

    def next_map_layer(self, reverse=False):
        self.current_map_layer += -1 if reverse else 1
        self.current_map_layer %= len(self.map_layer_names)
        selected = self.map_layer_names[self.current_map_layer]
        self.hud.notification('LayerMap selected: %s' % selected)

    def load_map_layer(self, unload=False):
        selected = self.map_layer_names[self.current_map_layer]
        if unload:
            self.hud.notification('Unloading map layer: %s' % selected)
            self.world.unload_map_layer(selected)
        else:
            self.hud.notification('Loading map layer: %s' % selected)
            self.world.load_map_layer(selected)

    # def toggle_radar(self):
    #     if self.radar_sensor is None:
    #         self.radar_sensor = RadarSensor(self.player)
    #     elif self.radar_sensor.sensor is not None:
    #         self.radar_sensor.sensor.destroy()
    #         self.radar_sensor = None

    def modify_vehicle_physics(self, actor):
        # If actor is not a vehicle, we cannot use the physics control
        try:
            physics_control = actor.get_physics_control()
            physics_control.use_sweep_wheel_collision = True
            actor.apply_physics_control(physics_control)
        except Exception:
            pass

    def tick(self, clock):
        self.hud.tick(self, clock)

    def render(self, display):
        self.camera_manager.render(display)
        self.hud.render(display)

    def destroy_sensors(self):
        self.camera_manager.sensor.destroy()
        self.camera_manager.sensor = None
        self.camera_manager.index = None

    def destroy(self):
        # if self.radar_sensor is not None:
        #     self.toggle_radar()
        sensors = [
            self.camera_manager.sensor,
            self.collision_sensor.sensor,
            self.lane_invasion_sensor.sensor,
            self.gnss_sensor.sensor,
            self.imu_sensor.sensor,
        ]
        for sensor in sensors:
            if sensor is not None:
                sensor.stop()
                sensor.destroy()
        if self.player is not None:
            self.player.destroy()


def sensor_callback(sensor_data, sensor_queue, sensor_name):
    if 'lidar' in sensor_name:
        sensor_data.save_to_disk(
            os.path.join(
                '../outputs/output_synchronized', '%06d.ply' % sensor_data.frame
            )
        )
    if 'camera' in sensor_name:
        sensor_data.save_to_disk(
            os.path.join(
                '../outputs/output_synchronized', '%06d.png' % sensor_data.frame
            )
        )
    sensor_queue.put((sensor_data.frame, sensor_name))


def main():
    actor_list = []
    sensor_list = []
    try:
        # First of all, we need to create the client that will send the requests, assume port is 2000
        client = carla.Client('localhost', 2000)
        client.set_timeout(2.0)

        # Retrieve the world that is currently running
        world = client.get_world()
        # world = client.load_world('Town02') # you can also retrive another world by specifically defining
        blueprint_library = world.get_blueprint_library()
        # Set weather for your world
        # weather = carla.WeatherParameters(
        #     cloudiness=10.0, precipitation=20.0, fog_density=10.0
        # )
        # world.set_weather(weather)

        # set synchorinized mode
        original_settings = world.get_settings()
        settings = world.get_settings()
        settings.fixed_delta_seconds = 0.05
        settings.synchronous_mode = True
        world.apply_settings(settings)

        traffic_manager = client.get_trafficmanager()
        traffic_manager.set_synchronous_mode(True)

        # create sensor queue
        sensor_queue = Queue()

        # create the ego vehicle
        vehicles_blueprints = world.get_blueprint_library().filter('*vehicle*')
        ego_vehicle_bp = random.choice(vehicles_blueprints)
        # black color
        ego_vehicle_bp.set_attribute('color', '0, 0, 0')
        # get a random valid occupation in the world
        transform = random.choice(world.get_map().get_spawn_points())
        # spawn the vehilce
        ego_vehicle = world.spawn_actor(ego_vehicle_bp, transform)
        # set the vehicle autopilot mode
        # ego_vehicle.set_autopilot(True)
        # collect all actors to destroy when we quit the script
        actor_list.append(ego_vehicle)

        # create directory for outputs
        output_path = '../outputs/output_synchronized'
        if not os.path.exists(output_path):
            os.makedirs(output_path)

        # 下面要添加传感器,这里以添加RGB相机为例:
        # Create a transform to place the camera on top of the vehicle

        # We create the camera through a blueprint that defines its properties
        camera_bp = world.get_blueprint_library().find('sensor.camera.rgb')

        # get the blueprint for this sensor
        # change the dimensions of the image
        camera_bp.set_attribute('image_size_x', f'{1920}')  # 图像宽度
        camera_bp.set_attribute('image_size_y', f'{1080}')  # 图像高度
        camera_bp.set_attribute('fov', '120')  # 水平视长角 （度）
        camera_bp.set_attribute('sensor_tick', '0.5')  # 消息间隔时间 （s）
        camera_bp.set_attribute('shutter_speed', '1/1000')
        camera_bp.set_attribute('iso', '100000')

        # add a camera
        # camera relative position related to the vehicle
        camera_transform = carla.Transform(carla.Location(x=1.5, z=2.4))
        camera = world.spawn_actor(camera_bp, camera_transform, attach_to=ego_vehicle)
        # set the callback function
        camera.listen(lambda image: sensor_callback(image, sensor_queue, "camera"))
        sensor_list.append(camera)

        # # we also add a lidar on it
        # lidar_bp = blueprint_library.find('sensor.lidar.ray_cast')
        # lidar_bp.set_attribute('channels', str(32))
        # lidar_bp.set_attribute('points_per_second', str(90000))
        # lidar_bp.set_attribute('rotation_frequency', str(40))
        # lidar_bp.set_attribute('range', str(20))

        # # set the relative location
        # lidar_location = carla.Location(0, 0, 2)
        # lidar_rotation = carla.Rotation(0, 0, 0)
        # lidar_transform = carla.Transform(lidar_location, lidar_rotation)
        # # spawn the lidar
        # lidar = world.spawn_actor(lidar_bp, lidar_transform, attach_to=ego_vehicle)
        # lidar.listen(
        #     lambda point_cloud: sensor_callback(point_cloud, sensor_queue, "lidar")
        # )
        # sensor_list.append(lidar)

        pygame.init()
        pygame.font.init()
        # 使用pygame创建显示窗口
        # 设置显示窗口大小
        windos_width = 1280
        windos_height = 720
        display = pygame.display.set_mode(
            (windos_width, windos_height), pygame.HWSURFACE | pygame.DOUBLEBUF
        )
        display.fill((0, 0, 0))
        pygame.display.flip()
        # 创建车辆信息打印窗口
        hud = HUD(windos_width, windos_height)
        # 创建世界变量控制器，用于控制天气、地图等信息(禁用了车辆信息更换)
        # 创建键盘控制器，用于操控车辆
        world_controler = World(world, hud, ego_vehicle)
        controller = KeyboardControl(world_controler, False)

        while True:
            world.tick()
            # set the sectator to follow the ego vehicle
            spectator = world.get_spectator()
            transform = ego_vehicle.get_transform()
            spectator.set_transform(
                carla.Transform(
                    transform.location + carla.Location(z=20), carla.Rotation(pitch=-90)
                )
            )
            clock = pygame.time.Clock()
            clock.tick_busy_loop(20)
            if controller.parse_events(client, world_controler, clock, True):
                break
            world_controler.tick(clock)
            world_controler.render(display)
            pygame.display.flip()
            # As the queue is blocking, we will wait in the queue.get() methods
            # until all the information is processed and we continue with the next frame.
            try:
                for i in range(0, len(sensor_list)):
                    s_frame = sensor_queue.get(True, 0.05)
                    print("    Frame: %d   Sensor: %s" % (s_frame[0], s_frame[1]))

            except Empty:
                print("   Some of the sensor information is missed")

    finally:
        world.apply_settings(original_settings)
        print('destroying actors')
        client.apply_batch([carla.command.DestroyActor(x) for x in actor_list])
        for sensor in sensor_list:
            sensor.destroy()
        print('done.')

        if world_controler is not None:
            world_controler.destroy()
        pygame.quit()
        settings = world.get_settings()
        settings.synchronous_mode = False
        settings.fixed_delta_seconds = None
        world.apply_settings(settings)
        print("finally processed!")


if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        print(' - Exited by user.')
